"""
@File    :   rsa_encryption.py
@Time    :   2024/02/11 11:55:54
@Author  :   Qrj
@Email   :   923207736@qq.com
@description   :   rsa非对称加解密
"""

import base64
from functools import lru_cache

import rsa
from asyncer import asyncify

from src.utils.path import PathUtil


class RsaEncryptionUtil:
    """Rsa加解密对象"""

    @staticmethod
    def generate_rsa(key_len: int = 2048):
        # 获取rsa加解密秘钥的文件夹路径
        rsa_encryption_path = PathUtil.get_rsa_encryption_path()
        if not rsa_encryption_path.exists():
            # 文件夹不存在，创建文件夹
            rsa_encryption_path.mkdir(parents=True, exist_ok=True)
        # 获取公钥与私钥文件位置
        rsa_public_key_file_path = PathUtil.get_rsa_public_key_file_path()
        rsa_private_key_file_path = PathUtil.get_rsa_private_key_file_path()
        # 判断公钥或私钥文件是否存在
        if rsa_public_key_file_path.exists() or rsa_private_key_file_path.exists():
            raise Exception("公钥或私钥文件已存在")

        # 创建公钥与私钥
        public_key, private_key = rsa.newkeys(key_len)
        # 保存公钥文件
        with open(rsa_public_key_file_path.as_posix(), "w") as public_f:
            public_f.write(public_key.save_pkcs1().decode())
        # 保存私钥文件
        with open(rsa_private_key_file_path.as_posix(), "w") as private_f:
            private_f.write(private_key.save_pkcs1().decode())

    @staticmethod
    @lru_cache
    def get_public_key() -> rsa.PublicKey:
        rsa_public_key_file_path = PathUtil.get_rsa_public_key_file_path()
        if not rsa_public_key_file_path.exists():
            # 当公钥文件不存在时, 试图去生成公钥文件
            RsaEncryptionUtil.generate_rsa()

        # 读取rsa公钥文件
        with open(rsa_public_key_file_path.as_posix(), "r") as f:
            public_key_obj = rsa.PublicKey.load_pkcs1(f.read().encode())
        return public_key_obj

    @staticmethod
    @lru_cache
    def get_private_key() -> rsa.PrivateKey:
        rsa_rivate_key_file_path = PathUtil.get_rsa_private_key_file_path()
        if not rsa_rivate_key_file_path.exists():
            # 当私钥文件不存在时, 试图去生成私钥文件
            RsaEncryptionUtil.generate_rsa()

        # 读取rsa私钥文件
        with open(rsa_rivate_key_file_path.as_posix(), "r") as f:
            private_key_obj = rsa.PrivateKey.load_pkcs1(f.read().encode())
        return private_key_obj

    @staticmethod
    def rsa_encrypt(encrypt_str: str, auto_decode: bool = False) -> bytes | str:
        """Rsa加密后使用base64进行编码

        Args:
            encrypt_str (str): 需加密的内容
            auto_decode: 是否自动解码

        Returns:
            bytes: 加密生成的bytes或经解码后的str
        """
        # 获取公钥进行加密
        public_key = RsaEncryptionUtil.get_public_key()
        encrypt_data = rsa.encrypt(encrypt_str.encode(), public_key)
        # base64编码
        encrypt_data = base64.b64encode(encrypt_data)
        if auto_decode:
            encrypt_data = encrypt_data.decode()
        return encrypt_data

    @staticmethod
    def rsa_decrypt(decrypt_data: str | bytes) -> str:
        """Rsa解密

        Args:
            decrypt_data (str | bytes): 需解密的数据

        Returns:
            str: 解密后的数据
        """
        # 获取私钥进行解密
        private_key = RsaEncryptionUtil.get_private_key()
        if isinstance(decrypt_data, str):
            # 将解密数据统一进行编码为bytes
            decrypt_data = decrypt_data.encode()
        # base64解码
        decrypt_data = base64.b64decode(decrypt_data)
        decrypt_result = rsa.decrypt(decrypt_data, private_key).decode()
        return decrypt_result

    @staticmethod
    async def async_rsa_encrypt(
        encrypt_str: str, auto_decode: bool = False
    ) -> bytes | str:
        """Rsa加密后使用base64进行编码(异步)

        Args:
            encrypt_str (str): 需加密的内容
            auto_decode: 是否自动解码

        Returns:
            bytes: 加密生成的bytes或经解码后的str
        """
        return await asyncify(RsaEncryptionUtil.rsa_encrypt)(
            encrypt_str=encrypt_str, auto_decode=auto_decode
        )

    @staticmethod
    async def async_rsa_decrypt(decrypt_data: str | bytes) -> str:
        """Rsa解密

        Args:
            decrypt_data (str | bytes): 需解密的数据

        Returns:
            str: 解密后的数据
        """
        return await asyncify(RsaEncryptionUtil.rsa_decrypt)(decrypt_data=decrypt_data)
